import os
from distutils.util import execute

from langchain.chains.sql_database.query import create_sql_query_chain
from langchain_community.tools import QuerySQLDatabaseTool
from langchain_community.utilities import SQLDatabase
from langchain_core.prompts import PromptTemplate
from langchain_core.runnables import RunnablePassthrough
from langchain_openai import ChatOpenAI
from oauthlib.uri_validate import query

# 配置http和https的服务器代理
os.environ["http_proxy"] = "127.0.0.1:8888"
os.environ["https_proxy"] ="127.0.0.1:8888"
os.environ["OPENAI_API_KEY"] = "你的密钥"
# 创建OpenAI模型
model = ChatOpenAI()
# 创建数据库url
url = "mysql+pymysql://root:123456@127.0.0.1:3306/你的数据库"
# 创建数据库对象
db = SQLDatabase.from_uri(url)
# 数据库和模型结合
sql_chain = create_sql_query_chain(model, db)
# 模型根据问题生成sql语句
execute_query = sql_chain.invoke({"question": "请查找用户admin的信息"})
print(execute_query)
# 创建数据库执行sql语句工具
execute_sql_tool = QuerySQLDatabaseTool(db=db)
# 执行sql语句, 返回结果
resp = execute_sql_tool.invoke(execute_query)
print(resp)
# 创建提示模板
answer_prompt = PromptTemplate.from_template(
    """给定以下用户问题，SQL执行的结果，回答问题
    Question:{question}
    Execute_Query:{execute_query}
    """
)
# 创建回答链
answer_chain = answer_prompt|model
# 模型根据问题和sql结果生成回答
print(answer_chain.invoke({"question": "请查找用户admin的信息", "execute_query": resp}).content)

